#!/usr/bin/python

#import os
import sys
import numpy as np
import matplotlib.pyplot as plt

"""
from optparse import OptionParser
parser = OptionParser("usage:  %prog [options] arg1 ")
parser.add_option("-k", "--kmcolumn", dest="kmcolumn")
parser.add_option("-t", "--transient", dest="transient")
"""

# FIXME: Distribution is working!!
def distMaxwellBoltzmann(KT):
  # Rejection method
  while(True):
    # 1. Generate an exponential random variable with lambda = 1
    Y  = np.random.exponential()
    # 2. Generate an uniform random variable
    U = np.random.uniform()
    # 3. Evaluate f(x)/(c*g(x))
    W = 0.5*(Y**2)*np.exp(-(Y**2)*0.5 +Y)
    if U <= W :
      return np.sqrt(KT)*Y
    else :
      pass


# Particle class
class Particle:
  def setX(self,x):
    self.x_new = x
  def setY(self,y):
    self.y_new = y
  def setVx(self,vx):
    self.vx_new = vx
  def setVy(self,vy):
    self.vy_new = vy
  def setSpeed(self):
    self.speed_new = 0.0
  def save_old(self):
    self.x_old  = self.x_new
    self.y_old  = self.y_new
    self.vx_old = self.vx_new
    self.vy_old = self.vy_new
    self.speed_old = self.speed_new

# Gas of particles class
# define boundary conditions

class Gas2D:
  # Box init dimensions, N
  def __init__(self, xlen, ylen, N, KT):
    self.N = N
    self.xlen = xlen
    self.ylen = ylen
    self.particle = [ Particle() for i in range(self.N)]
    self.KT = KT
    # uniform random positions
    for pi in self.particle:
      pi.x_new = np.random.uniform(0, self.xlen)
      pi.y_new = np.random.uniform(0, self.ylen) 
    # TODO: maxwell boltzmann velocities distribution
    for pi in self.particle:
      theta = np.random.uniform(0,1)
      pi.speed_new = distMaxwellBoltzmann(self.KT)
      pi.vx_new = pi.speed_new*np.cos(theta*2*np.pi)
      pi.vy_new = pi.speed_new*np.sin(theta*2*np.pi)
      
  def nextStep(self, dt):
    for pi in self.particle:
      pi.save_old()
      pi.x_new = pi.x_old + dt*pi.vx_old
      pi.y_new = pi.y_old + dt*pi.vy_old
              
      delta_x =  self.xlen - pi.x_new
      delta_y =  self.ylen - pi.y_new
      # TODO: in this case we don't care about the interactions.
      if pi.x_new < 0:
        pi.x_new  = -pi.x_new
        pi.vx_new = -pi.vx_new
      elif delta_x < 0 :
        # A Wall collision occurs in x axis
        pi.x_new  = self.xlen + delta_x
        pi.vx_new = -pi.vx_new
      if pi.y_new < 0:
        pi.y_new  = -pi.y_new
        pi.vy_new = -pi.vy_new
      elif delta_y < 0:
        # A Wall collision occurs in x axis
        pi.y_new  = self.ylen + delta_y
        pi.vy_new = -pi.vy_new
      

      

# Is a 2-D box with lenght xL yL
if __name__ == "__main__":
  # And God says...
  N  = 1000
  KT_eq = 10
  gas   = Gas2D(10, 1, N, KT_eq)
  gasX  = [gas.particle[i].x_new for i in range(gas.N)]
  gasY  = [gas.particle[i].y_new for i in range(gas.N)]
  gasVX = [gas.particle[i].vx_new for i in range(gas.N)]
  gasVY = [gas.particle[i].vy_new for i in range(gas.N)]

  # ---- BEGIN FIGURE 1
  plt.figure(1)
  outputfilename = 'IdealGas_NpartBoxIniCond.pdf'
  font = {'family' : 'normal',
        'weight' : 'bold',
        'size'   : 24}
  plt.rc('font', **font)
  plt.rc('text', usetex=True)
  plt.axes().set_aspect('equal')
  plt.yticks(np.arange(0, 2, 0.5))

  ax = plt.gca()
  ax.set_xlim([0,gas.xlen])
  ax.set_ylim([0,gas.ylen])
  ax.quiver(gasX, gasY, gasVX, gasVY, color='red')
  plt.savefig(outputfilename,transparent=True)
  # END FIGURE 1

  
  dt = 0.001
  ttime  = 10000
  gasX_t = np.zeros(ttime)
  gasY_t = np.zeros(ttime)
  for i in range(ttime):
    gas.nextStep(dt)
    gasX_t[i] = gas.particle[1].x_new
    gasY_t[i] = gas.particle[1].y_new
  
  print len(gasX_t), gasX_t
  # print len(ttime), ttime
  
  # ---- BEGIN FIGURE  2
  # Make a plot
  plt.figure(2)
  outputfilename = 'IdealGas_1-particleBox.pdf'
  font = {'family' : 'normal',
        'weight' : 'bold',
        'size'   : 24}
  plt.rc('font', **font)
  plt.rc('text', usetex=True)
  plt.axes().set_aspect('equal')
  plt.yticks(np.arange(0, 2, 0.5))
  plt.xlim((0,10))
  plt.ylim((0,1))
  
  plt.plot(gasX_t, gasY_t)
  plt.savefig(outputfilename,transparent=True)
  # END FIGURE 2
  
  # ---- BEGIN FIGURE 3
  plt.figure(3)
  outputfilename = 'IdealGasVelocityDist.pdf'
  plt.xlabel('$v$')
  plt.ylabel('$n$')
  plt.hist([gas.particle[i].speed_new for i in range(gas.N)])
  plt.savefig(outputfilename,transparent=True)
  
  # END FIGURE 3
  


